package org.sction.security.shiro;

import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.Collection;
import java.util.HashSet;
import java.util.Iterator;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;

import javax.sql.DataSource;

import org.apache.log4j.Logger;
import org.apache.shiro.authc.AuthenticationException;
import org.apache.shiro.authc.AuthenticationInfo;
import org.apache.shiro.authc.AuthenticationToken;
import org.apache.shiro.authc.SimpleAuthenticationInfo;
import org.apache.shiro.authz.AuthorizationException;
import org.apache.shiro.authz.AuthorizationInfo;
import org.apache.shiro.authz.SimpleAuthorizationInfo;
import org.apache.shiro.cas.CasAuthenticationException;
import org.apache.shiro.cas.CasToken;
import org.apache.shiro.realm.AuthorizingRealm;
import org.apache.shiro.subject.PrincipalCollection;
import org.apache.shiro.subject.SimplePrincipalCollection;
import org.apache.shiro.util.CollectionUtils;
import org.apache.shiro.util.JdbcUtils;
import org.apache.shiro.util.StringUtils;
import org.jasig.cas.client.authentication.AttributePrincipal;
import org.jasig.cas.client.validation.Assertion;
import org.jasig.cas.client.validation.Cas20ServiceTicketValidator;
import org.jasig.cas.client.validation.Saml11TicketValidator;
import org.jasig.cas.client.validation.TicketValidationException;
import org.jasig.cas.client.validation.TicketValidator;
import org.springframework.jdbc.core.JdbcTemplate;

/**
 * @author 作者 : sorc
 * @version 创建时间：2012-12-25 上午8:28:59
 * 
 *          类说明
 */

public class ShiroCasRealm extends AuthorizingRealm {
	private static final Logger logger = Logger.getLogger(ShiroCasRealm.class);

	protected DataSource dataSource;
	protected String userRolesQuery;
	protected String permissionsQuery;
	protected boolean permissionsLookupEnabled = true;

	public static final String DEFAULT_REMEMBER_ME_ATTRIBUTE_NAME = "longTermAuthenticationRequestTokenUsed";
	public static final String DEFAULT_VALIDATION_PROTOCOL = "CAS";

	protected String casServerUrlPrefix;
	protected String casService;
	protected String validationProtocol;
	protected String rememberMeAttributeName;

	private TicketValidator ticketValidator;

	protected String defaultRoles;
	protected String defaultPermissions;

	protected JdbcTemplate jdbc;
	protected String sql;

	protected String roleAttributeNames;
	protected String permissionAttributeNames;

	Set<String> roleNames;
	Set<String> permissions;

	public ShiroCasRealm() {
		validationProtocol = DEFAULT_VALIDATION_PROTOCOL;
		rememberMeAttributeName = "longTermAuthenticationRequestTokenUsed";
		setAuthenticationTokenClass(CasToken.class);
	}

	/**
	 * 初始化数据
	 */
	protected void onInit() {
		super.onInit();
		ensureTicketValidator();
	}

	/**
	 * 确认令牌有效性
	 * 
	 * @return
	 */
	protected TicketValidator ensureTicketValidator() {
		if (ticketValidator == null)
			ticketValidator = createTicketValidator();
		logger.info("ticketValidator:" + ticketValidator.toString());
		return ticketValidator;
	}

	/**
	 * 创建令牌验证器
	 * 
	 * @return
	 */
	protected TicketValidator createTicketValidator() {
		String urlPrefix = getCasServerUrlPrefix();
		logger.info("ValidationProtocol:" + getValidationProtocol() + urlPrefix);
		if ("saml".equalsIgnoreCase(getValidationProtocol()))
			return new Saml11TicketValidator(urlPrefix);
		else
			return new Cas20ServiceTicketValidator(urlPrefix);
	}

	/**
	 * 认证
	 */
	protected AuthenticationInfo doGetAuthenticationInfo(
			AuthenticationToken token) throws AuthenticationException {
		CasToken casToken;
		String ticket;
		TicketValidator ticketValidator;
		casToken = (CasToken) token;
		if (token == null) {
			return null;
		}

		ticket = (String) casToken.getCredentials();
		if (!StringUtils.hasText(ticket)) {
			return null;
		} else {
			try {
				ticketValidator = ensureTicketValidator();
				PrincipalCollection principalCollection;
				Assertion casAssertion = ticketValidator.validate(ticket,
						getCasService());
				AttributePrincipal casPrincipal = casAssertion.getPrincipal();
				String userId = casPrincipal.getName();
				Map<?, ?> attributes = casPrincipal.getAttributes();
				logger.info("Validate ticket : {" + ticket
						+ "} in CAS server : {" + getCasServerUrlPrefix()
						+ "} to retrieve user : {" + userId + "}");
				casToken.setUserId(userId);
				String rememberMeAttributeName = getRememberMeAttributeName();
				String rememberMeStringValue = (String) attributes
						.get(rememberMeAttributeName);
				boolean isRemembered = rememberMeStringValue != null
						&& Boolean.parseBoolean(rememberMeStringValue);
				if (isRemembered)
					casToken.setRememberMe(true);
				List<Object> principals = CollectionUtils.asList(new Object[] {
						userId, attributes });

				principalCollection = new SimplePrincipalCollection(principals,
						getName());

				return new SimpleAuthenticationInfo(principalCollection, ticket);

			} catch (TicketValidationException e) {
				logger.error(e.getMessage(), e);
			}
		}
		throw new CasAuthenticationException((new StringBuilder())
				.append("Unable to validate ticket [").append(ticket)
				.append("]").toString());
	}

	/**
	 * 授权
	 */
	protected AuthorizationInfo doGetAuthorizationInfo(
			PrincipalCollection principals) {

		if (principals == null)
			throw new AuthorizationException(
					"PrincipalCollection method argument cannot be null.");
		Connection conn = null;

		roleNames = split(defaultRoles);
		permissions = split(defaultPermissions);

		SimpleAuthorizationInfo info = new SimpleAuthorizationInfo();
		String userid = (String) getAvailablePrincipal(principals);

		if (dataSource != null) {
			try {
				conn = dataSource.getConnection();
				roleNames = getRoleNamesForUser(conn, userid);
				if (permissionsLookupEnabled)
					permissions = getPermissions(conn, userid, roleNames);
			} catch (SQLException e) {
				String message = (new StringBuilder())
						.append("There was a SQL error while authorizing user [")
						.append(userid).append("]").toString();

				logger.error(message, e);
				throw new AuthorizationException(message, e);
			} finally {
				JdbcUtils.closeConnection(conn);
			}
			JdbcUtils.closeConnection(conn);
			info.setRoles(roleNames);
			info.setStringPermissions(permissions);
			return info;
		} else {
			if (jdbc != null) {
				if (sql == null) {
					throw new AuthorizationException("请设置sql属性");
				}
				sql = sql.replaceAll("\\?", "'" + userid + "'");
				try {
					@SuppressWarnings("unchecked")
					List<Map<String, Object>> list = jdbc.queryForList(sql);
					logger.info("来自本地数据库的用户属性:" + list);
					if (list != null) {
						for (Map<String, Object> map : list) {
							putRolesPermissions(map);
						}
					}
				} catch (Exception e) {
					logger.error(sql, e);
				}
			} else {
				@SuppressWarnings("unchecked")
				Map<String, Object> map = (Map<String, Object>) principals
						.asList().get(1);
				logger.info("来自CAS的用户属性:" + map);
				putRolesPermissions(map);
			}
		}
		return info;
	}

	/**
	 * 将map数据取出并填充到用户的角色和权限
	 * 
	 * @param map
	 */
	private void putRolesPermissions(Map<String, Object> map) {
		String roles;
		String pops;
		if (roleAttributeNames != null && !roleAttributeNames.trim().equals("")) {
			Set<String> attributeNames = split(roleAttributeNames);
			for (Iterator<String> i$ = attributeNames.iterator(); i$.hasNext(); roleNames
					.addAll(split(roles))) {
				String attributeName = i$.next();
				roles = (String) map.get(attributeName);
			}
		}
		if (permissionAttributeNames != null
				&& !permissionAttributeNames.trim().equals("")) {

			Set<String> attributeNames = split(permissionAttributeNames);
			for (Iterator<String> i$ = attributeNames.iterator(); i$.hasNext(); permissions
					.addAll(split(pops))) {
				String attributeName = i$.next();
				pops = (String) map.get(attributeName);
			}
		}
	}

	/**
	 * 分割用，划分的字符串为Set集合
	 * 
	 * @param s
	 * @return
	 */
	private Set<String> split(String s) {
		Set<String> set = new HashSet<String>();
		String elements[] = StringUtils.split(s, ',');
		if (elements != null) {
			for (int i$ = 0; i$ < elements.length; i$++) {
				String element = elements[i$];
				if (StringUtils.hasText(element))
					set.add(element.trim());
			}
		}
		return set;
	}

	/**
	 * 通过数据库查询用户角色
	 * 
	 * @param conn
	 * @param username
	 * @return
	 * @throws SQLException
	 */
	protected Set<String> getRoleNamesForUser(Connection conn, String username)
			throws SQLException {
		PreparedStatement ps;
		ResultSet rs;
		Set<String> roleNames;
		ps = null;
		rs = null;
		roleNames = new LinkedHashSet<String>();
		ps = conn.prepareStatement(userRolesQuery);
		ps.setString(1, username);
		rs = ps.executeQuery();
		while (rs.next()) {
			String roleName = rs.getString(1);
			if (roleName != null)
				roleNames.add(roleName);
			else
				logger.warn((new StringBuilder())
						.append("Null role name found while retrieving role names for user [")
						.append(username).append("]").toString());
		}
		JdbcUtils.closeResultSet(rs);
		JdbcUtils.closeStatement(ps);
		return roleNames;
	}

	/**
	 * 通过数据库角色查询权限
	 * 
	 * @param conn
	 * @param username
	 * @param roleNames
	 * @return
	 * @throws SQLException
	 */
	protected Set<String> getPermissions(Connection conn, String username,
			Collection<String> roleNames) throws SQLException {
		PreparedStatement ps = null;
		Set<String> permissions = new LinkedHashSet<String>();
		ResultSet rs = null;
		ps = conn.prepareStatement(permissionsQuery);
		Iterator<String> i$ = roleNames.iterator();
		while (i$.hasNext()) {
			String roleName = (String) i$.next();
			ps.setString(1, roleName);
			String permissionString;
			for (rs = ps.executeQuery(); rs.next(); permissions
					.add(permissionString))
				permissionString = rs.getString(1);
			JdbcUtils.closeResultSet(rs);
		}
		JdbcUtils.closeStatement(ps);
		return permissions;
	}

	public String getCasServerUrlPrefix() {
		return casServerUrlPrefix;
	}

	public void setCasServerUrlPrefix(String casServerUrlPrefix) {
		this.casServerUrlPrefix = casServerUrlPrefix;
	}

	public String getCasService() {
		return casService;
	}

	public void setCasService(String casService) {
		this.casService = casService;
	}

	public String getValidationProtocol() {
		return validationProtocol;
	}

	public void setValidationProtocol(String validationProtocol) {
		this.validationProtocol = validationProtocol;
	}

	public String getRememberMeAttributeName() {
		return rememberMeAttributeName;
	}

	public void setRememberMeAttributeName(String rememberMeAttributeName) {
		this.rememberMeAttributeName = rememberMeAttributeName;
	}

	public String getDefaultRoles() {
		return defaultRoles;
	}

	public void setDefaultRoles(String defaultRoles) {
		this.defaultRoles = defaultRoles;
	}

	public void setDefaultPermissions(String defaultPermissions) {
		this.defaultPermissions = defaultPermissions;
	}

	public void setRoleAttributeNames(String roleAttributeNames) {
		this.roleAttributeNames = roleAttributeNames;
	}

	public void setPermissionAttributeNames(String permissionAttributeNames) {
		this.permissionAttributeNames = permissionAttributeNames;
	}

	public void setSql(String sql) {
		this.sql = sql;
	}

	public void setJdbc(JdbcTemplate jdbc) {
		this.jdbc = jdbc;
	}

	public void setDataSource(DataSource dataSource) {
		this.dataSource = dataSource;
	}

	public void setUserRolesQuery(String userRolesQuery) {
		this.userRolesQuery = userRolesQuery;
	}

	public void setPermissionsQuery(String permissionsQuery) {
		this.permissionsQuery = permissionsQuery;
	}

	public void setPermissionsLookupEnabled(boolean permissionsLookupEnabled) {
		this.permissionsLookupEnabled = permissionsLookupEnabled;
	}

}
